import asyncio
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Form
from utils.msg_redis import MessageRedis
from utils.msg_websocket import ConnectionManager
from utils.model_infer import detect_aoi_lot, detect_pmi_folder, detect_aoi_wafer
import json
from router.api import api_router
from utils.post_item import AOILotInferItem, PMILotInferItem, AOIWaferInferItem
from concurrent.futures import ThreadPoolExecutor, ProcessPoolExecutor
from fastapi.staticfiles import StaticFiles
from fastapi import UploadFile
from typing import List
# from tempfile import NamedTemporaryFile
import shutil
from utils.gen_vtest import gen_aoi
import os
from pathlib import Path
from mmengine.registry import Registry
# from utils.celery_task import celery_app, detect_aoi_lot_celery
root_path = os.path.dirname(os.path.abspath(__file__))

from weice_package.cls_infer import cls_init
# import det_infer
# import seg_infer
from weice_package.det_infer_multi_batch import get_det_model, get_det_model_pmi
from weice_package.seg_infer_multi_batch import get_seg_model


app = FastAPI()

app.include_router(api_router, prefix="/api")

# =============aoi====================
config_file_det = "gpu_e2e/configs/yolox_l_1024_4xb12-100e_flip_aoi_weice_copy.py"
checkpoint_file_det = "gpu_e2e/float_models/best_coco_bbox_mAP_epoch_119.pth"

config_path_seg = "gpu_e2e/configs/fpn_r18_256x256_weice_laketi_b8_nmpad_1009.py"
checkpoint_path_seg = "gpu_e2e/float_models/iter_320000_aoi_nmpad_1009.pth"

config_file_cls = "gpu_e2e/configs/resnet50_bs32_100e_lr01_pretrained_weice_pad_4cls_v7_0925.py"
checkpoint_file_cls = "gpu_e2e/float_models/best_accuracy_top1_epoch_55.pth"


# ===================pmi===================
pmi_config_file_det = "gpu_e2e/configs/yolox_m_640_2xb8-128e_flip_pmi_weice.py"
pmi_checkpoint_file_det = "gpu_e2e/float_models/best_coco_bbox_mAP_epoch_123_pmi_1007.pth"
pmi_config_path_seg = "gpu_e2e/configs/fpn_r18_128x128_weice_pmi_b8_no_extend.py"
pmi_checkpoint_path_seg = "gpu_e2e/float_models/iter_320000_pmi_1018.pth"


# 初始化 AOI 模型
# with Registry('scope').switch_scope_and_registry('mmdet'):
#     det_model = det_infer.det_init(config_file_det, checkpoint_file_det)
# with Registry('scope').switch_scope_and_registry('mmseg'):
#     seg_model = seg_infer.seg_init(config_path_seg, checkpoint_path_seg)
with Registry('scope').switch_scope_and_registry('mmcls'):    
    cls_model = cls_init(config_file_cls, checkpoint_file_cls)

# 初始化 AOI 模型（多batch）
multi_batch_det_model = get_det_model(checkpoint_file_det)
multi_batch_seg_model = get_seg_model(checkpoint_path_seg)

# 初始化 PMI 模型（多batch）
pmi_det_model = get_det_model_pmi(pmi_checkpoint_file_det)
pmi_seg_model = get_seg_model(pmi_checkpoint_path_seg)


loop = asyncio.get_event_loop()
executor = ThreadPoolExecutor()

manager = ConnectionManager()
redis = MessageRedis()

app.mount("/static", StaticFiles(directory="static"), name="static")


@app.websocket("/msg/{client}")
async def websocket_endpoint(client: str, websocket: WebSocket):
    await manager.connect(client, websocket)
    try:
        while True:
            # 接收客户端发送的消息
            data = await websocket.receive_text()
            await manager.send_text(client, data)
    except WebSocketDisconnect:
        manager.disconnect(client, websocket)


# aoi 自动分析
@app.post('/api/model/auto')
async def aoi_auto(item: AOILotInferItem):
    # task_status = await redis.get_message("task")

    # if task_status == "start":
    #     return {"code": 1, "msg": "后台任务已启动, 请等待任务完成"}

    lot_folder = "/data/lot/test_lot"
    distance = item.distance
    set_num_zhenhen = item.set_num_zhenhen
    set_diagonal_length = item.set_diagonal_length
    area_scale = item.area_scale

    # 污染面积
    set_area = item.set_area

    # 针痕面积百分比
    set_per_area = float(item.set_per_area/100)

    await loop.run_in_executor(None, detect_aoi_lot, multi_batch_det_model, multi_batch_seg_model, cls_model, lot_folder, distance, set_num_zhenhen, set_diagonal_length, set_area, set_per_area, area_scale, manager)

    # 多进程运行, 接口直接返回
    # print("多线程运行")
    # executor.submit(detect_aoi_lot, aoi_yolox_model, aoi_seg_model, cls_model, lot_folder, distance, set_num_zhenhen, set_diagonal_length, set_area, set_per_area)

    # 同步运行
    # detect_aoi_lot(multi_batch_det_model, multi_batch_seg_model, cls_model, lot_folder, distance, set_num_zhenhen, set_diagonal_length, set_area, set_per_area, area_scale, manager)

    # 使用 celery 运行耗时任务
    # task = detect_aoi_lot_celery.delay(aoi_yolox_model, aoi_seg_model, cls_model, lot_folder, distance, set_num_zhenhen, set_diagonal_length, set_area, set_per_area)

    # # 适用 python 3.9+
    # asyncio.to_thread(detect_aoi_lot, lot_folder, distance, set_num_zhenhen, set_diagonal_length, set_area, set_per_area)

    return {"code": 0, "msg": "开始自动推理"}

# pmi 自动分析
@app.post('/api/pmi/auto')
async def pmi_auto(item: PMILotInferItem):
    # task_status = await redis.get_message("task")

    # print("task_status: ", task_status)

    # if task_status == "start":
    #     return {"code": 1, "msg": "后台任务已启动, 请等待任务完成"}

    pmi_folder = "/data/pmi/"
    device_id = item.device_id
    distance = item.distance
    area_scale = item.area_scale

    # 针痕面积百分比
    set_per_area = float(item.set_per_area/100) 

    # detect_pmi_folder(pmi_det_model, pmi_seg_model, device_id, pmi_folder, distance, set_per_area, area_scale, manager)

    await loop.run_in_executor(None, detect_pmi_folder, pmi_det_model, pmi_seg_model, device_id, pmi_folder, distance, set_per_area, area_scale, manager)

    return {"code": 0, "msg": "开始自动推理 %s 的 Failed 图片 " % device_id}


# aoi 上传图片
@app.post('/api/aoi/upload')
async def img_upload(images: List[UploadFile]):

    aoi_dir = f"{root_path}/static/aoi"

    shutil.rmtree(aoi_dir)
    os.mkdir(aoi_dir)

    all_temp_filename = []

    for file in images:
        try:
            suffix = Path(file.filename).suffix[1:]
            # with NamedTemporaryFile(delete=False, suffix=suffix, dir=aoi_dir) as tmp:
            with open(f"{aoi_dir}/{file.filename}", "wb") as tmp:
                # shutil.copyfileobj(file.file, tmp)
                content = await(file.read())

                if suffix in ["jpg", "jpeg"]:
                    # JPG 和 JPEG 文件结尾标识缺失
                    if not content.endswith(b'\xff\xd9'):
                        content += b'\xff\xd9'
                elif suffix in ["png"]:
                    # png 文件结尾标识缺失
                    if not content.endswith(b'\xaeB`\x82'):
                        content += b'\xaeB`\x82'
                tmp.write(content)
                tmp_file_name = Path(tmp.name).name
                all_temp_filename.append(tmp_file_name)
        finally:
            file.file.close()

    gen_aoi(aoi_dir)

    return {"msg": "上传图片成功", "code": 0}


# aoi 手动分析
@app.post('/api/aoi/img-infer')
async def img_infer(item: AOIWaferInferItem):

    task_status = await redis.get_message("task")

    if task_status == "start":
        return {"code": 1, "msg": "后台任务已启动, 请等待任务完成"} 

    distance = item.distance
    set_num_zhenhen = item.set_num_zhenhen
    set_diagonal_length = item.set_diagonal_length
    area_scale = item.area_scale

    # 污染面积
    set_area = item.set_area

    # 针痕面积百分比
    set_per_area = float(item.set_per_area/100)

    aoi_dir = f"{root_path}/static/aoi"

    # loop = asyncio.get_running_loop()

    # result = await loop.run_in_executor(None, detect_aoi_wafer, aoi_dir, distance, set_num_zhenhen, set_diagonal_length, set_area, set_per_area)

    result = detect_aoi_wafer(multi_batch_det_model, multi_batch_seg_model, cls_model, aoi_dir, distance, set_num_zhenhen, set_diagonal_length, set_area, set_per_area, area_scale)

    return result


@app.get('/')
def index():
    return "首页"